{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "(vta-mat-mult-opt)=\n",
        "# 分块矩阵乘法\n",
        "\n",
        "\n",
        "**原作者**: [Thierry Moreau](https://homes.cs.washington.edu/~moreau/)\n",
        "\n",
        "本教程概述了如何在 VTA 设计中使用 TVM 有效地映射矩阵乘法。建议先学习 {ref}`basic-mat-mult` 教程。\n",
        "\n",
        "在本教程中，将演示 TVM 调度优化，将大型神经网络算子分解为较小的块，以在有限的硬件加速器资源内实现计算。\n",
        "\n",
        "## RPC 设置\n",
        "\n",
        "首先编程 Pynq 的 FPGA 并构建它的 RPC 运行时。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "import set_env"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import tvm\n",
        "from tvm import te\n",
        "import vta\n",
        "import numpy as np\n",
        "from tvm import rpc\n",
        "from tvm.contrib import utils\n",
        "from vta.testing import simulator\n",
        "\n",
        "# Load VTA parameters from the 3rdparty/vta-hw/config/vta_config.json file\n",
        "env = vta.get_env()\n",
        "\n",
        "# We read the Pynq RPC host IP address and port number from the OS environment\n",
        "host = os.environ.get(\"VTA_RPC_HOST\", \"192.168.2.99\")\n",
        "port = int(os.environ.get(\"VTA_RPC_PORT\", \"9091\"))\n",
        "\n",
        "# We configure both the bitstream and the runtime system on the Pynq\n",
        "# to match the VTA configuration specified by the vta_config.json file.\n",
        "if env.TARGET == \"pynq\":\n",
        "\n",
        "    # Make sure that TVM was compiled with RPC=1\n",
        "    assert tvm.runtime.enabled(\"rpc\")\n",
        "    remote = rpc.connect(host, port)\n",
        "\n",
        "    # Reconfigure the JIT runtime\n",
        "    vta.reconfig_runtime(remote)\n",
        "\n",
        "    # Program the FPGA with a pre-compiled VTA bitstream.\n",
        "    # You can program the FPGA with your own custom bitstream\n",
        "    # by passing the path to the bitstream file instead of None.\n",
        "    vta.program_fpga(remote, bitstream=None)\n",
        "\n",
        "# In simulation mode, host the RPC server locally.\n",
        "elif env.TARGET in [\"sim\", \"tsim\"]:\n",
        "    remote = rpc.LocalSession()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 声明计算\n",
        "\n",
        "作为第一步，需要描述矩阵乘法的计算。将矩阵乘法定义为全连接层中的计算，由其 batch size、输入通道和输出通道定义。它们必须是 VTA 张量形状的整数倍：`BATCH`、`BLOCK_IN` 和 `BLOCK_OUT`。\n",
        "\n",
        "在矩阵乘法中添加额外的算子，这些算子对输出进行了移位（shifting）和剪切（clipping），以模拟定点矩阵乘法，然后是修正的线性激活。将全连通层的 TVM 数据流图描述如下：\n",
        "\n",
        "```{image} images/fc_dataflow.png\n",
        ":align: center\n",
        "```\n",
        "\n",
        "此计算被故意设置得太大，以至于不能一次全部放入 VTA 的 on-chip buffer。因此，在调度阶段，将依靠计算分块策略将计算分解为可管理的块。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# 全连接层 1024 x 1024\n",
        "batch_size = 1\n",
        "in_channels = 1024\n",
        "out_channels = 1024\n",
        "num_ops = in_channels * out_channels * batch_size * 2\n",
        "\n",
        "assert batch_size % env.BATCH == 0\n",
        "assert in_channels % env.BLOCK_IN == 0\n",
        "assert out_channels % env.BLOCK_OUT == 0\n",
        "\n",
        "# 推导出平铺的张量形状\n",
        "data_shape = (\n",
        "    batch_size // env.BATCH, \n",
        "    in_channels // env.BLOCK_IN,\n",
        "    env.BATCH, env.BLOCK_IN\n",
        ")\n",
        "weight_shape = (\n",
        "    out_channels // env.BLOCK_OUT,\n",
        "    in_channels // env.BLOCK_IN,\n",
        "    env.BLOCK_OUT,\n",
        "    env.BLOCK_IN,\n",
        ")\n",
        "output_shape = (\n",
        "    batch_size // env.BATCH, \n",
        "    out_channels // env.BLOCK_OUT, \n",
        "    env.BATCH, env.BLOCK_OUT\n",
        ")\n",
        "\n",
        "# Reduction axes\n",
        "ic = te.reduce_axis((0, in_channels // env.BLOCK_IN), name=\"ic\")\n",
        "ic_tns = te.reduce_axis((0, env.BLOCK_IN), name=\"ic_tns\")\n",
        "\n",
        "# Input placeholder tensors\n",
        "data = te.placeholder(data_shape, name=\"data\", dtype=env.inp_dtype)\n",
        "weight = te.placeholder(weight_shape, name=\"weight\", dtype=env.wgt_dtype)\n",
        "\n",
        "# Copy buffers\n",
        "data_buf = te.compute(data_shape, lambda *i: data(*i), \"data_buf\")\n",
        "weight_buf = te.compute(weight_shape, lambda *i: weight(*i), \"weight_buf\")\n",
        "\n",
        "# 声明矩阵乘法计算\n",
        "res_gemm = te.compute(\n",
        "    output_shape,\n",
        "    lambda bo, co, bi, ci: te.sum(\n",
        "        data_buf[bo, ic, bi, ic_tns].astype(env.acc_dtype)\n",
        "        * weight_buf[co, ic, ci, ic_tns].astype(env.acc_dtype),\n",
        "        axis=[ic, ic_tns],\n",
        "    ),\n",
        "    name=\"res_gem\",\n",
        ")\n",
        "\n",
        "# 为定点归一化（fix-point normalization）添加 shift stage\n",
        "res_shr = te.compute(output_shape, lambda *i: res_gemm(*i) >> env.INP_WIDTH, name=\"res_shr\")\n",
        "\n",
        "# 将值裁剪到 (0, input max value)\n",
        "inp_max = (1 << (env.INP_WIDTH - 1)) - 1\n",
        "res_max = te.compute(output_shape, lambda *i: tvm.te.max(res_shr(*i), 0), \"res_max\")\n",
        "res_min = te.compute(output_shape, lambda *i: tvm.te.min(res_max(*i), inp_max), \"res_min\")\n",
        "\n",
        "# 在返回结果之前，对输入数据类型应用类型转换\n",
        "res = te.compute(output_shape, lambda *i: res_min(*i).astype(env.inp_dtype), name=\"res\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 调度计算\n",
        "\n",
        "查看一组必要的调度变换，以有效的方式将矩阵乘法映射到 VTA。这些包括：\n",
        "\n",
        "- 分块计算（Computation blocking）\n",
        "- Lowering 到 VTA 硬件 intrinsics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), weight: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), res: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)):\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>), <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        data_buf <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>allocate([<span style=\"color: #008000\">1024</span>], <span style=\"color: #BA2121\">&quot;int8&quot;</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
              "        weight_buf <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>allocate([<span style=\"color: #008000\">1048576</span>], <span style=\"color: #BA2121\">&quot;int8&quot;</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
              "        res_gem <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>allocate([<span style=\"color: #008000\">1024</span>], <span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
              "        data_buf_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>data_buf)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1, i3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_1: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3\n",
              "            data_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>data<span style=\"color: #A2F; font-weight: bold\">.</span>data)\n",
              "            data_buf_1[cse_var_1] <span style=\"color: #A2F; font-weight: bold\">=</span> data_1[cse_var_1]\n",
              "        weight_buf_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1048576</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>weight_buf)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1, i2, i3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_2: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i0 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16384</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3\n",
              "            weight_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1048576</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>weight<span style=\"color: #A2F; font-weight: bold\">.</span>data)\n",
              "            weight_buf_1[cse_var_2] <span style=\"color: #A2F; font-weight: bold\">=</span> weight_1[cse_var_2]\n",
              "        res_gem_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_gem)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> co, ci <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>):\n",
              "            res_gem_1[co <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ci] <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #008000\">0</span>\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> ic, ic_tns <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_3: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> co <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ci\n",
              "                res_gem_1[cse_var_3] <span style=\"color: #A2F; font-weight: bold\">=</span> res_gem_1[cse_var_3] <span style=\"color: #A2F; font-weight: bold\">+</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, data_buf_1[ic <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic_tns]) <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, weight_buf_1[co <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16384</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ci <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic_tns])\n",
              "        res_gem_2 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_gem)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1, i3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_4: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3\n",
              "            res_gem_2[cse_var_4] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>shift_right(res_gem_1[cse_var_4], <span style=\"color: #008000\">8</span>)\n",
              "        res_gem_3 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_gem)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1, i3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_5: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3\n",
              "            res_gem_3[cse_var_5] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>max(res_gem_2[cse_var_5], <span style=\"color: #008000\">0</span>)\n",
              "        res_gem_4 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_gem)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1, i3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_6: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3\n",
              "            res_gem_4[cse_var_6] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>min(res_gem_3[cse_var_6], <span style=\"color: #008000\">127</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1, i3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_7: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3\n",
              "            res_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res<span style=\"color: #A2F; font-weight: bold\">.</span>data)\n",
              "            res_1[cse_var_7] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int8&quot;</span>, res_gem_4[cse_var_7])\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# 创建 TVM 调度\n",
        "s = te.create_schedule(res.op)\n",
        "# 查看默认调度\n",
        "tvm.lower(s, [data, weight, res], simple_mode=True).show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 分块计算\n",
        "\n",
        "在默认情况下，矩阵乘法对于激活或权重来说太大了，无法一次性适应 VTA 的 on-chip buffer。将 (1, 1024)×(1024, 1024) 矩阵乘法分成更小的 (1, 256) × (256, 256) 矩阵乘法，这样中间张量就可以装进加速器的 on-chip SRAM 中。这种方法类似于将分块技术应用于 CPU 和 GPU，以提高缓存命中率（cache hit rate）。\n",
        "\n",
        "沿着每个轴执行分块（batch 轴不受影响，因为正在执行单 batch 推理）。也保持最内侧的 tensorization 轴不变，以便 TVM 能够进行模式匹配的 tensorization。在下面的图表中展示了分块在计算调度上的结果：\n",
        "\n",
        "```{image} images/blocking.png\n",
        ":align: center\n",
        ":width: 480px\n",
        "```\n",
        "\n",
        "````{admonition} 循环分割（splitting）和重新排序（reordering）后的代码等价于下面的伪代码。忽略 batch 轴，因为在这个例子中只执行单 batch 推断：\n",
        ":class: alert alert-info\n",
        "```c\n",
        "for (int oc_out = 0; oc_out < 4; ++oc_out) {\n",
        "  // Initialization loop\n",
        "  for (int oc_inn = 0; oc_inn < 16; ++oc_inn) {\n",
        "   for (int oc_tns = 0; oc_tns < 16; ++oc_tns) {\n",
        "    int j = (oc_out * 16 + oc_inn) * 16 + oc_tns;\n",
        "    C[0][j] = 0;\n",
        "   }\n",
        "  }\n",
        "  for (int ic_out = 0; ic_out < 4; ++ic_out) {\n",
        "   // Block loop\n",
        "   for (int oc_inn = 0; oc_inn < 16; ++oc_inn) {\n",
        "    for (int ic_inn = 0; ic_inn < 16; ++ic_inn) {\n",
        "     // Tensorization loop\n",
        "     for (int oc_tns = 0; oc_tns < 16; ++oc_tns) {\n",
        "      for (int ic_tns = 0; ic_tns < 16; ++ic_tns) {\n",
        "       int i = (ic_out * 16 + ic_inn) * 16 + ic_tns;\n",
        "       int j = (oc_out * 16 + oc_inn) * 16 + oc_tns;\n",
        "       C[0][i] = C[0][i] + A[0][i] * B[j][i];\n",
        "      }\n",
        "     }\n",
        "    }\n",
        "   }\n",
        "  }\n",
        " }\n",
        "}\n",
        "```\n",
        "````"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), weight: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), res: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)):\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>), <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        data_buf <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>allocate([<span style=\"color: #008000\">1024</span>], <span style=\"color: #BA2121\">&quot;int8&quot;</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
              "        weight_buf <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>allocate([<span style=\"color: #008000\">1048576</span>], <span style=\"color: #BA2121\">&quot;int8&quot;</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
              "        res_gem <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>allocate([<span style=\"color: #008000\">256</span>], <span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
              "        data_buf_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>data_buf)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1, i3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_1: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3\n",
              "            data_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>data<span style=\"color: #A2F; font-weight: bold\">.</span>data)\n",
              "            data_buf_1[cse_var_1] <span style=\"color: #A2F; font-weight: bold\">=</span> data_1[cse_var_1]\n",
              "        weight_buf_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1048576</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>weight_buf)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1, i2, i3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_2: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i0 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16384</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3\n",
              "            weight_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1048576</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>weight<span style=\"color: #A2F; font-weight: bold\">.</span>data)\n",
              "            weight_buf_1[cse_var_2] <span style=\"color: #A2F; font-weight: bold\">=</span> weight_1[cse_var_2]\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1_outer <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">4</span>):\n",
              "            res_gem_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">256</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_gem)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> co_init, ci_init <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "                res_gem_1[co_init <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ci_init] <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #008000\">0</span>\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> ic_outer, co, ic_inner, ci, ic_tns <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_3: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> co <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ci\n",
              "                res_gem_1[cse_var_3] <span style=\"color: #A2F; font-weight: bold\">=</span> res_gem_1[cse_var_3] <span style=\"color: #A2F; font-weight: bold\">+</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, data_buf_1[ic_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic_inner <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic_tns]) <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, weight_buf_1[i1_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">262144</span> <span style=\"color: #A2F; font-weight: bold\">+</span> co <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16384</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">4096</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic_inner <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ci <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic_tns])\n",
              "            res_gem_2 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">256</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_gem)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> i1, i3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_4: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3\n",
              "                res_gem_2[cse_var_4] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>shift_right(res_gem_1[cse_var_4], <span style=\"color: #008000\">8</span>)\n",
              "            res_gem_3 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">256</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_gem)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> i1, i3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_5: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3\n",
              "                res_gem_3[cse_var_5] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>max(res_gem_2[cse_var_5], <span style=\"color: #008000\">0</span>)\n",
              "            res_gem_4 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">256</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_gem)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> i1, i3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_6: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3\n",
              "                res_gem_4[cse_var_6] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>min(res_gem_3[cse_var_6], <span style=\"color: #008000\">127</span>)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> i1_inner, i3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_7: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1_inner <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span>\n",
              "                res_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res<span style=\"color: #A2F; font-weight: bold\">.</span>data)\n",
              "                res_1[i1_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_7 <span style=\"color: #A2F; font-weight: bold\">+</span> i3] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int8&quot;</span>, res_gem_4[cse_var_7 <span style=\"color: #A2F; font-weight: bold\">+</span> i3])\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# Let's define tiling sizes (expressed in multiples of VTA tensor shape size)\n",
        "b_block = 1 // env.BATCH\n",
        "i_block = 256 // env.BLOCK_IN\n",
        "o_block = 256 // env.BLOCK_OUT\n",
        "\n",
        "# Tile the output tensor along the batch and output channel dimensions\n",
        "# (since by default we are doing single batch inference, the split along\n",
        "#  the batch dimension has no effect)\n",
        "b, oc, b_tns, oc_tns = s[res].op.axis\n",
        "b_out, b_inn = s[res].split(b, b_block)\n",
        "oc_out, oc_inn = s[res].split(oc, o_block)\n",
        "s[res].reorder(b_out, oc_out, b_inn, oc_inn)\n",
        "\n",
        "# Move intermediate computation into each output compute tile\n",
        "s[res_gemm].compute_at(s[res], oc_out)\n",
        "s[res_shr].compute_at(s[res], oc_out)\n",
        "s[res_max].compute_at(s[res], oc_out)\n",
        "s[res_min].compute_at(s[res], oc_out)\n",
        "\n",
        "# Apply additional loop split along reduction axis (input channel)\n",
        "b_inn, oc_inn, b_tns, oc_tns = s[res_gemm].op.axis\n",
        "ic_out, ic_inn = s[res_gemm].split(ic, i_block)\n",
        "\n",
        "# Reorder axes. We move the ic_out axis all the way out of the GEMM\n",
        "# loop to block along the reduction axis\n",
        "s[res_gemm].reorder(ic_out, b_inn, oc_inn, ic_inn, b_tns, oc_tns, ic_tns)\n",
        "\n",
        "# Let's look at the current TVM schedule after blocking\n",
        "tvm.lower(s, [data, weight, res], simple_mode=True).show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### lowering 复制到 DMA 传输\n",
        "\n",
        "接下来，将 buffer 作用域设置为相应的 on-chip VTA SRAM buffer。将 load 循环移动到矩阵乘法计算循环中，以使它们适合于 on-chip SRAM buffer。最后，用 DMA 复制实用程序对 load/store 循环外轴进行注解，以在 VTA 上执行批量内存传输。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Set scope of SRAM buffers\n",
        "s[data_buf].set_scope(env.inp_scope)\n",
        "s[weight_buf].set_scope(env.wgt_scope)\n",
        "s[res_gemm].set_scope(env.acc_scope)\n",
        "s[res_shr].set_scope(env.acc_scope)\n",
        "s[res_min].set_scope(env.acc_scope)\n",
        "s[res_max].set_scope(env.acc_scope)\n",
        "\n",
        "# Block data and weight cache reads\n",
        "s[data_buf].compute_at(s[res_gemm], ic_out)\n",
        "s[weight_buf].compute_at(s[res_gemm], ic_out)\n",
        "\n",
        "# Use DMA copy pragma on DRAM->SRAM operations\n",
        "s[data_buf].pragma(s[data_buf].op.axis[0], env.dma_copy)\n",
        "s[weight_buf].pragma(s[weight_buf].op.axis[0], env.dma_copy)\n",
        "\n",
        "# Use DMA copy pragma on SRAM->DRAM operation\n",
        "# (this implies that these copies should be performed along b_inn,\n",
        "# or result axis 2)\n",
        "s[res].pragma(s[res].op.axis[2], env.dma_copy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Lowering 计算到 VTA Compute Intrinsics\n",
        "\n",
        "最后阶段是通过将矩阵乘法映射到张量 intrinsics，将 shift 映射到矢量 ALU，从而将计算循环 lowering 到 VTA 硬件 intrinsics。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.command_handle\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.command_handle\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.command_handle\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_sync\n",
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), weight: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), res: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)):\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>), <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1_outer <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">4</span>):\n",
              "            vta <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int32()\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_scope&quot;</span>, <span style=\"color: #008000\">2</span>):\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>)\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_uop_scope&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAPushGEMMOp&quot;</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>uop_push(<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> ic_outer <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">4</span>):\n",
              "                cse_var_1: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> ic_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span>\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_scope&quot;</span>, <span style=\"color: #008000\">1</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTALoadBuffer2D&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>tvm_thread_context(T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>command_handle()), data<span style=\"color: #A2F; font-weight: bold\">.</span>data, cse_var_1, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">2</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTALoadBuffer2D&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>tvm_thread_context(T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>command_handle()), weight<span style=\"color: #A2F; font-weight: bold\">.</span>data, i1_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">1024</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_1, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>)\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_scope&quot;</span>, <span style=\"color: #008000\">2</span>)\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>)\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_uop_scope&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAPushGEMMOp&quot;</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">16</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>uop_push(<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>)\n",
              "            T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>)\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_scope&quot;</span>, <span style=\"color: #008000\">2</span>):\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_uop_scope&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAPushALUOp&quot;</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>uop_push(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_uop_scope&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAPushALUOp&quot;</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>uop_push(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_uop_scope&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAPushALUOp&quot;</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>uop_push(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">127</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>)\n",
              "            T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_scope&quot;</span>, <span style=\"color: #008000\">3</span>)\n",
              "            T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> i1_inner <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">16</span>):\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAStoreBuffer2D&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>tvm_thread_context(T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>command_handle()), i1_inner, <span style=\"color: #008000\">4</span>, res<span style=\"color: #A2F; font-weight: bold\">.</span>data, i1_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i1_inner, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>)\n",
              "            T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>)\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_sync()\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>)\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# Apply tensorization over the batch tensor tile axis\n",
        "s[res_gemm].tensorize(b_tns, env.gemm)\n",
        "\n",
        "# Add an ALU pragma over the shift and clipping operations\n",
        "s[res_shr].pragma(s[res_shr].op.axis[0], env.alu)\n",
        "s[res_min].pragma(s[res_min].op.axis[0], env.alu)\n",
        "s[res_max].pragma(s[res_max].op.axis[0], env.alu)\n",
        "\n",
        "# Let's look at the final lowered TVM schedule after lowering memory\n",
        "# loads/stores down to DMA copy intrinsics, and the computation down to\n",
        "# VTA compute intrinsics.\n",
        "vta.lower(s, [data, weight, res], simple_mode=True).show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## TVM 计算和验证\n",
        "\n",
        "在指定调度之后，可以将其编译为 TVM 函数。保存模块，这样就可以通过 RPC 发送它。运行该函数并对 numpy 实现进行验证，以确保其正确性。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "[20:12:08] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
            "2025-07-27 20:12:08.741 INFO load_module /tmp/tmplc94wcpg/gemm.o\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Execution statistics:\n",
            "\tinp_load_nbytes :             4096\n",
            "\twgt_load_nbytes :          1048576\n",
            "\tacc_load_nbytes :                0\n",
            "\tuop_load_nbytes :               20\n",
            "\tout_store_nbytes:             1024\n",
            "\tgemm_counter    :             4096\n",
            "\talu_counter     :              192\n",
            "Successful blocked matrix multiply test!\n"
          ]
        }
      ],
      "source": [
        "# Compile the TVM module\n",
        "my_gemm = vta.build(\n",
        "    s, [data, weight, res], tvm.target.Target(\"ext_dev\", host=env.target_host), name=\"my_gemm\"\n",
        ")\n",
        "temp = utils.tempdir()\n",
        "my_gemm.save(temp.relpath(\"gemm.o\"))\n",
        "remote.upload(temp.relpath(\"gemm.o\"))\n",
        "f = remote.load_module(\"gemm.o\")\n",
        "\n",
        "# Get the remote device context\n",
        "ctx = remote.ext_dev(0)\n",
        "\n",
        "# Initialize the data and weight arrays randomly in the int range of (-128, 128]\n",
        "data_np = np.random.randint(-128, 128, size=(batch_size, in_channels)).astype(data.dtype)\n",
        "weight_np = np.random.randint(-128, 128, size=(out_channels, in_channels)).astype(weight.dtype)\n",
        "\n",
        "# Apply packing to the data and weight arrays from a 2D to a 4D packed layout\n",
        "data_packed = data_np.reshape(\n",
        "    batch_size // env.BATCH, env.BATCH, in_channels // env.BLOCK_IN, env.BLOCK_IN\n",
        ").transpose((0, 2, 1, 3))\n",
        "weight_packed = weight_np.reshape(\n",
        "    out_channels // env.BLOCK_OUT, env.BLOCK_OUT, in_channels // env.BLOCK_IN, env.BLOCK_IN\n",
        ").transpose((0, 2, 1, 3))\n",
        "\n",
        "# Format the input/output arrays with tvm.nd.array to the DLPack standard\n",
        "data_nd = tvm.nd.array(data_packed, ctx)\n",
        "weight_nd = tvm.nd.array(weight_packed, ctx)\n",
        "res_nd = tvm.nd.array(np.zeros(output_shape).astype(res.dtype), ctx)\n",
        "\n",
        "# Clear stats\n",
        "if env.TARGET in [\"sim\", \"tsim\"]:\n",
        "    simulator.clear_stats()\n",
        "\n",
        "# Invoke the module to perform the computation\n",
        "f(data_nd, weight_nd, res_nd)\n",
        "\n",
        "# Verify against numpy implementation\n",
        "res_ref = np.dot(data_np.astype(env.acc_dtype), weight_np.T.astype(env.acc_dtype))\n",
        "res_ref = res_ref >> env.INP_WIDTH\n",
        "res_ref = np.clip(res_ref, 0, inp_max)\n",
        "res_ref = res_ref.astype(res.dtype)\n",
        "res_ref = res_ref.reshape(\n",
        "    batch_size // env.BATCH, env.BATCH, out_channels // env.BLOCK_OUT, env.BLOCK_OUT\n",
        ").transpose((0, 2, 1, 3))\n",
        "np.testing.assert_equal(res_ref, res_nd.numpy())\n",
        "\n",
        "# Print stats\n",
        "if env.TARGET in [\"sim\", \"tsim\"]:\n",
        "    sim_stats = simulator.stats()\n",
        "    print(\"Execution statistics:\")\n",
        "    for k, v in sim_stats.items():\n",
        "        print(\"\\t{:<16}: {:>16}\".format(k, v))\n",
        "\n",
        "print(\"Successful blocked matrix multiply test!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 小结\n",
        "\n",
        "```{topic} 小结\n",
        "本教程演示了 TVM 调度原语如何为矩阵乘法示例实现分块计算。这允许将任意大的计算映射到有限的硬件加速器资源上。\n",
        "```"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
